﻿
// DiskWriterDlg.cpp: 实现文件
//

#include "pch.h"
#include "framework.h"
#include "DiskWriter.h"
#include "DiskWriterDlg.h"
#include "afxdialogex.h"
#include "Disk.h"

#ifdef _DEBUG
#define new DEBUG_NEW
#endif


// 用于应用程序“关于”菜单项的 CAboutDlg 对话框

class CAboutDlg : public CDialogEx
{
public:
	CAboutDlg();

// 对话框数据
#ifdef AFX_DESIGN_TIME
	enum { IDD = IDD_ABOUTBOX };
#endif

	protected:
	virtual void DoDataExchange(CDataExchange* pDX);    // DDX/DDV 支持

// 实现
protected:
	DECLARE_MESSAGE_MAP()
};

CAboutDlg::CAboutDlg() : CDialogEx(IDD_ABOUTBOX)
{
}

void CAboutDlg::DoDataExchange(CDataExchange* pDX)
{
	CDialogEx::DoDataExchange(pDX);
}

BEGIN_MESSAGE_MAP(CAboutDlg, CDialogEx)
END_MESSAGE_MAP()


// CDiskWriterDlg 对话框



CDiskWriterDlg::CDiskWriterDlg(CWnd* pParent /*=nullptr*/)
	: CDialogEx(IDD_DISKWRITER_DIALOG, pParent)
	, m_szFilePath(_T(""))
{
	m_hIcon = AfxGetApp()->LoadIcon(IDR_MAINFRAME);
}

void CDiskWriterDlg::DoDataExchange(CDataExchange* pDX)
{
	CDialogEx::DoDataExchange(pDX);
	DDX_Text(pDX, IDC_EDIT_FILE, m_szFilePath);
	DDX_Control(pDX, IDC_BUTTON_BROWSE, m_buttonBrowse);
	DDX_Control(pDX, IDC_COMBO_DEVICE, m_comboDevice);
	DDX_Control(pDX, IDC_PROGRESS, m_progress);
	DDX_Control(pDX, IDC_COMBO_FILE_CHUNK_SIZE, m_comboChunkSize);
	DDX_Control(pDX, IDC_BUTTON_WRITE, m_buttonWrite);
}

BEGIN_MESSAGE_MAP(CDiskWriterDlg, CDialogEx)
	ON_WM_SYSCOMMAND()
	ON_WM_PAINT()
	ON_WM_QUERYDRAGICON()
	ON_WM_DEVICECHANGE()
	ON_MESSAGE(WM_USER_UPDATE_UI, OnUpdateUI)
	ON_BN_CLICKED(IDC_BUTTON_BROWSE, &CDiskWriterDlg::OnBnClickedButtonBrowse)
	ON_BN_CLICKED(IDC_BUTTON_WRITE, &CDiskWriterDlg::OnBnClickedButtonWrite)
END_MESSAGE_MAP()


// CDiskWriterDlg 消息处理程序

BOOL CDiskWriterDlg::OnInitDialog()
{
	CDialogEx::OnInitDialog();
	ASSERT((IDM_ABOUTBOX & 0xFFF0) == IDM_ABOUTBOX);
	ASSERT(IDM_ABOUTBOX < 0xF000);

	CMenu* pSysMenu = GetSystemMenu(FALSE);
	if (pSysMenu != nullptr)
	{
		BOOL bNameValid;
		CString strAboutMenu;
		bNameValid = strAboutMenu.LoadString(IDS_ABOUTBOX);
		ASSERT(bNameValid);
		if (!strAboutMenu.IsEmpty())
		{
			pSysMenu->AppendMenu(MF_SEPARATOR);
			pSysMenu->AppendMenu(MF_STRING, IDM_ABOUTBOX, strAboutMenu);
		}
	}
	SetIcon(m_hIcon, TRUE);
	SetIcon(m_hIcon, FALSE);

	//Initialize controls
	HBITMAP hBitmap = LoadBitmap(AfxGetInstanceHandle(), MAKEINTRESOURCE(IDB_BITMAP1));
	m_buttonBrowse.SetBitmap(hBitmap);

	// Device list.
	m_comboDevice.SetCurSel(0);

	// Writing progressbar.
	m_progress.SetRange32(0, 100);

	// File chunk size list.
	m_comboChunkSize.AddString(TEXT("128B"));
	m_comboChunkSize.AddString(TEXT("256B"));
	m_comboChunkSize.AddString(TEXT("512B"));
	m_comboChunkSize.SetItemData(0, 128);
	m_comboChunkSize.SetItemData(1, 256);
	m_comboChunkSize.SetItemData(2, 512);
	m_comboChunkSize.SetCurSel(0);

	m_buttonWrite.EnableWindow(FALSE);

	// Register device notification
	DEV_BROADCAST_DEVICEINTERFACE NotificationFilter;
	ZeroMemory(&NotificationFilter, sizeof(NotificationFilter));
	NotificationFilter.dbcc_size = sizeof(DEV_BROADCAST_DEVICEINTERFACE);
	NotificationFilter.dbcc_devicetype = DBT_DEVTYP_DEVICEINTERFACE;

	HDEVNOTIFY hDevNotify = RegisterDeviceNotification(
		this->GetSafeHwnd(),
		&NotificationFilter,
		DEVICE_NOTIFY_WINDOW_HANDLE);

	if(!hDevNotify)
		AfxMessageBox(_T("Fail to register device change notification."));

	// Get logical drives.
	GetLogicalDevices();

	// 注销设备更改通知
	//UnregisterDeviceNotification(hDevNotify);

	return TRUE;  // 除非将焦点设置到控件，否则返回 TRUE
}

LRESULT CDiskWriterDlg::OnUpdateUI(WPARAM wParam, LPARAM lParam)
{
	switch(wParam)
	{
	case UI_SET_PROGRESS_RANGE:
		m_progress.SetRange32(0, (int)lParam);
	break;
	case UI_SET_PROGRESS_VALUE:
		m_progress.SetPos((int)lParam);
	break;
	case UI_WRITE_DISK_BEGIN:
		m_bWritting = TRUE;
		SetWriteButtonState();
	break;
	case UI_WRITE_DISK_FINISH:
		m_bWritting = FALSE;
		m_progress.SetPos(0);
		SetWriteButtonState();
		break;
	}
	return 0;
}

void CDiskWriterDlg::OnSysCommand(UINT nID, LPARAM lParam)
{
	if ((nID & 0xFFF0) == IDM_ABOUTBOX)
	{
		CAboutDlg dlgAbout;
		dlgAbout.DoModal();
	}
	else
	{
		CDialogEx::OnSysCommand(nID, lParam);
	}
}

// 如果向对话框添加最小化按钮，则需要下面的代码
//  来绘制该图标。  对于使用文档/视图模型的 MFC 应用程序，
//  这将由框架自动完成。

void CDiskWriterDlg::OnPaint()
{
	if (IsIconic())
	{
		CPaintDC dc(this); // 用于绘制的设备上下文

		SendMessage(WM_ICONERASEBKGND, reinterpret_cast<WPARAM>(dc.GetSafeHdc()), 0);

		// 使图标在工作区矩形中居中
		int cxIcon = GetSystemMetrics(SM_CXICON);
		int cyIcon = GetSystemMetrics(SM_CYICON);
		CRect rect;
		GetClientRect(&rect);
		int x = (rect.Width() - cxIcon + 1) / 2;
		int y = (rect.Height() - cyIcon + 1) / 2;

		// 绘制图标
		dc.DrawIcon(x, y, m_hIcon);
	}
	else
	{
		CDialogEx::OnPaint();
	}
}

//当用户拖动最小化窗口时系统调用此函数取得光标
//显示。
HCURSOR CDiskWriterDlg::OnQueryDragIcon()
{
	return static_cast<HCURSOR>(m_hIcon);
}

void CDiskWriterDlg::GetLogicalDevices()
{
	// GetLogicalDrives returns 0 on failure, or a bitmask representing
	// the drives available on the system (bit 0 = A:, bit 1 = B:, etc)
	unsigned long driveMask = GetLogicalDrives();
	int i = 0;
	ULONG pID;
	m_comboDevice.Clear();

	while(driveMask != 0)
	{
		if(driveMask & 1)
		{
			// the "A" in drivename will get incremented by the # of bits
			// we've shifted
			char drivename[] = "\\\\.\\A:\\";
			drivename[4] += i;
			if(CheckDriveType(drivename, &pID))
			{
				CString szItem;
				szItem.Format(TEXT("[%c:\\]"), drivename[4]);
				m_comboDevice.AddString(szItem);
				m_comboDevice.SetItemData(0, (unsigned long long)pID);
			}
		}
		driveMask >>= 1;
		m_comboDevice.SetCurSel(0);
		++i;
	}
}

void CDiskWriterDlg::SetWriteButtonState()
{
	UpdateData(TRUE);
	BOOL bFileSelected = !m_szFilePath.IsEmpty();
	BOOL bDeviceSelected = (m_comboDevice.GetCount() > 0);

	CFileStatus fs;
	BOOL bFileExists = CFile::GetStatus(m_szFilePath, fs);

	// set read and write buttons according to status of file/device
	m_buttonWrite.EnableWindow(!m_bWritting && bFileSelected && bDeviceSelected && bFileExists);
}

// support routine for winEvent - returns the drive letter for a given mask
//   taken from http://support.microsoft.com/kb/163503
char FirstDriveFromMask(ULONG unitmask)
{
	char i;

	for(i = 0; i < 26; ++i)
	{
		if(unitmask & 0x1)
		{
			break;
		}
		unitmask = unitmask >> 1;
	}

	return (i + 'A');
}

BOOL CDiskWriterDlg::OnDeviceChange(UINT nEventType, DWORD dwData)
{
	PDEV_BROADCAST_HDR lpdb = (PDEV_BROADCAST_HDR)dwData;
	switch(nEventType)
	{
	case DBT_DEVICEARRIVAL:
	{
		if(lpdb->dbch_devicetype == DBT_DEVTYP_VOLUME)
		{
			PDEV_BROADCAST_VOLUME lpdbv = (PDEV_BROADCAST_VOLUME)lpdb;
			if(DBTF_NET)
			{
				char ALET = FirstDriveFromMask(lpdbv->dbcv_unitmask);
				// add device to combo box (after sanity check that
				// it's not already there, which it shouldn't be)
				CString szItem;
				szItem.Format(TEXT("[%c:\\]"), ALET);
				if(m_comboDevice.FindString(0, szItem) == -1)
				{
					ULONG pID;
					char longname[] = "\\\\.\\A:\\";
					longname[4] = ALET;
					// checkDriveType gets the physicalID
					if(CheckDriveType(longname, &pID))
					{
						m_comboDevice.AddString(szItem);
						m_comboDevice.SetItemData(0, (unsigned long long)pID);
						if(m_comboDevice.GetCurSel() == -1)
							m_comboDevice.SetCurSel(0);

						SetWriteButtonState();
					}
				}
			}
		}
	}
	break;

	case DBT_DEVICEREMOVECOMPLETE:
	{
		if(lpdb->dbch_devicetype == DBT_DEVTYP_VOLUME)
		{
			PDEV_BROADCAST_VOLUME lpdbv = (PDEV_BROADCAST_VOLUME)lpdb;
			if(DBTF_NET)
			{
				char ALET = FirstDriveFromMask(lpdbv->dbcv_unitmask);
				//  find the device that was removed in the combo box,
				//  and remove it from there....
				//  "removeItem" ignores the request if the index is
				//  out of range, and findText returns -1 if the item isn't found.

				CString szItem;
				szItem.Format(TEXT("[%c:\\]"), ALET);
				UINT index = m_comboDevice.FindString(0, szItem);
				m_comboDevice.DeleteString(index);
				m_comboDevice.SetCurSel(0);
				SetWriteButtonState();
			}
		}
	}
	break;

	default:
		break;
	}

	return CDialogEx::OnDeviceChange(nEventType, dwData);
}

void OnWriteDisk(LPVOID param)
{
	Context* ctx = (Context*)param;

	HANDLE hVolume = GetHandleOnVolume(ctx->volumeID, GENERIC_WRITE);
	if(hVolume == INVALID_HANDLE_VALUE)
	{
		AfxGetMainWnd()->PostMessage(WM_USER_UPDATE_UI, UI_WRITE_DISK_FINISH, 0);
		return;
	}

	// Lock volume.
	DWORD deviceID = GetDeviceID(hVolume);
	if(!GetLockOnVolume(hVolume))
	{
		::CloseHandle(hVolume);

		AfxGetMainWnd()->PostMessage(WM_USER_UPDATE_UI, UI_WRITE_DISK_FINISH, 0);
		return;
	}

	// Unmount volume.
	if(!UnmountVolume(hVolume))
	{
		RemoveLockOnVolume(hVolume);
		::CloseHandle(hVolume);

		AfxGetMainWnd()->PostMessage(WM_USER_UPDATE_UI, UI_WRITE_DISK_FINISH, 0);
		return;
	}

	// Open Image file.
	HANDLE hFile = GetHandleOnFile(ctx->szFilePath, GENERIC_READ);
	if(hFile == INVALID_HANDLE_VALUE)
	{
		RemoveLockOnVolume(hVolume);
		::CloseHandle(hVolume);

		AfxGetMainWnd()->PostMessage(WM_USER_UPDATE_UI, UI_WRITE_DISK_FINISH, 0);
		return;
	}

	// Open device.
	HANDLE hRawDisk = GetHandleOnDevice(deviceID, GENERIC_WRITE);
	if(hRawDisk == INVALID_HANDLE_VALUE)
	{
		RemoveLockOnVolume(hVolume);
		::CloseHandle(hFile);
		::CloseHandle(hVolume);

		AfxGetMainWnd()->PostMessage(WM_USER_UPDATE_UI, UI_WRITE_DISK_FINISH, 0);
		return;
	}

	// Get sector info.
	unsigned long long nDiskSectorSize = 0;
	unsigned long long nDiskSectors = GetNumberOfSectors(hRawDisk, &nDiskSectorSize);
	if(nDiskSectors == 0)
	{
		RemoveLockOnVolume(hVolume);
		::CloseHandle(hRawDisk);
		::CloseHandle(hFile);
		::CloseHandle(hVolume);

		AfxGetMainWnd()->PostMessage(WM_USER_UPDATE_UI, UI_WRITE_DISK_FINISH, 0);
		return;
	}

	// Calculate the number of sectors occupied by file.
	unsigned long long nFileSectors = GetFileSizeInSectors(hFile, ctx->nFileChunkSize);
	if(nFileSectors == 0)
	{
		RemoveLockOnVolume(hVolume);
		::CloseHandle(hRawDisk);
		::CloseHandle(hFile);
		::CloseHandle(hVolume);

		AfxGetMainWnd()->PostMessage(WM_USER_UPDATE_UI, UI_WRITE_DISK_FINISH, 0);
		return;
	}

	// File size exceeds disk capacity.
	if(nFileSectors > nDiskSectors)
	{
		RemoveLockOnVolume(hVolume);
		::CloseHandle(hRawDisk);
		::CloseHandle(hFile);
		::CloseHandle(hVolume);

		CString szMsg;
		szMsg.Format(TEXT("File size exceeds disk capacity, "));
		::MessageBox(NULL, szMsg, TEXT("File Error"), MB_OK);

		AfxGetMainWnd()->PostMessage(WM_USER_UPDATE_UI, UI_WRITE_DISK_FINISH, 0);
		return;
	}

	// Set range of progressbar
	AfxGetMainWnd()->PostMessage(WM_USER_UPDATE_UI, UI_SET_PROGRESS_RANGE, (LPARAM)nFileSectors);
	AfxGetMainWnd()->PostMessage(WM_USER_UPDATE_UI, UI_SET_PROGRESS_VALUE, (LPARAM)0);
	AfxGetMainWnd()->PostMessage(WM_USER_UPDATE_UI, UI_WRITE_DISK_BEGIN, 0);

	// Write sectors to disk.
	BOOL bSuccess = TRUE;
	DWORD dwBytesRead = 0;
	DWORD dwBytesWrite = 0;
	char* data = new char[nDiskSectorSize];
	memset(data, 0, nDiskSectorSize);
	int nSectorWritten = 0;
	while(ReadFile(hFile, data, ctx->nFileChunkSize, &dwBytesRead, NULL) && dwBytesRead > 0)
	{
		if(WriteFile(hRawDisk, data, nDiskSectorSize, &dwBytesWrite, NULL) == FALSE || dwBytesWrite <= 0)
		{
			bSuccess = FALSE;
			break;
		}

		// Update progress.
		nSectorWritten++;
		AfxGetMainWnd()->PostMessage(WM_USER_UPDATE_UI, UI_SET_PROGRESS_VALUE, (LPARAM)nSectorWritten);
		memset(data, 0, nDiskSectorSize);
	}
	delete[] data;

	if(!bSuccess)
	{
		RemoveLockOnVolume(hVolume);
		::CloseHandle(hRawDisk);
		::CloseHandle(hFile);
		::CloseHandle(hVolume);

		CString szMsg = GetErrorMessage(TEXT("An error occurred when attempting to write data to disk."));
		::MessageBox(NULL, (LPCTSTR)szMsg, TEXT("Write Error"), MB_OK);

		AfxGetMainWnd()->PostMessage(WM_USER_UPDATE_UI, UI_WRITE_DISK_FINISH, 0);
		return;
	}

	RemoveLockOnVolume(hVolume);
	::CloseHandle(hRawDisk);
	::CloseHandle(hFile);
	::CloseHandle(hVolume);

	AfxGetMainWnd()->PostMessage(WM_USER_UPDATE_UI, UI_WRITE_DISK_FINISH, 0);
	::MessageBox(NULL, TEXT("Write Successful."), TEXT("Write Complete"), MB_OK);
}

void CDiskWriterDlg::OnBnClickedButtonBrowse()
{
	// TODO: 在此添加控件通知处理程序代码
	CFileDialog dialog(TRUE, NULL, NULL, OFN_FILEMUSTEXIST | OFN_PATHMUSTEXIST, _T("All Files (*.*)|*.*||"));
	if(dialog.DoModal() == IDOK)
	{
		m_szFilePath = dialog.GetPathName();
		UpdateData(FALSE);
		SetWriteButtonState();
	}
}


void CDiskWriterDlg::OnBnClickedButtonWrite()
{
	// TODO: 在此添加控件通知处理程序代码
	UpdateData(TRUE);

	// Is the file name empty?
	if(m_szFilePath.IsEmpty())
	{
		::MessageBox(NULL, TEXT("Please specify an image file to use."), TEXT("File Error"), MB_OK);
		return;
	}

	// Is it the same drive?
	CString szDriveName;
	m_comboDevice.GetLBText(m_comboDevice.GetCurSel(), szDriveName);
	if(m_szFilePath.GetAt(0) == szDriveName.GetAt(1))
	{
		::MessageBox(NULL, TEXT("Image file cannot be located on the target device."), TEXT("File Error"), MB_OK);
		return;
	}

	// File exists?
	CFileStatus fs;
	if(!CFile::GetStatus(m_szFilePath, fs))
	{
		::MessageBox(NULL, TEXT("The selected file does not exist."), TEXT("File Error"), MB_OK);
		return;
	}

	// Is file empty?
	if(fs.m_size == 0)
	{
		::MessageBox(NULL, TEXT("The specified file contains no data."), TEXT("File Error"), MB_OK);
		return;
	}

	// Build the drive letter as a const char* (without the surrounding brackets)
	szDriveName.Replace(TEXT("["), TEXT(""));
	szDriveName.Replace(TEXT("]"), TEXT(""));

	// Confirm if the operation should be continued.
	CString szWarning;
	szWarning.Format(
		TEXT("Writing to a physical device can corrupt the device.\n"
			"(Target Device: %s \"%s\")\n"
			"Are you sure you want to continue?"), szDriveName, GetDriveLabel(szDriveName));
	int result = ::MessageBox(NULL, szWarning, _T("Confirm Overwrite"), MB_YESNO | MB_ICONQUESTION);
	if(result == IDNO)
		return;

	// Get volume handle.
	int volumeID = szDriveName.GetAt(0) - 'A';

	// File chunk size.
	int nFileChunkSize = m_comboChunkSize.GetItemData(m_comboChunkSize.GetCurSel());

	// Begin thread.
	m_ctx.szFilePath = m_szFilePath;
	m_ctx.volumeID = volumeID;
	m_ctx.nFileChunkSize = nFileChunkSize;
	AfxBeginThread((AFX_THREADPROC)OnWriteDisk, (LPVOID)&m_ctx);
}